#!/bin/sh
#######################################################################################
#
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# security-tool licensed under the Mulan PSL v1.
# You can use this software according to the terms and conditions of the Mulan PSL v1.
# You may obtain a copy of Mulan PSL v1 at:
#     http://license.coscl.org.cn/MulanPSL
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Mulan PSL v1 for more details.
# Description: Configuration file for the ima mesurement tags.
#
#######################################################################################

# ima policy generated patterns
# These patterns are used to identify the start and end of the generated area in the IMA policy file.
readonly IMA_POLICY_GENERATED_PATTERN_START='# Generated by security-tool'
readonly IMA_POLICY_GENERATED_PATTERN_END='# End of generated by security-tool'
# IMA policy file path
readonly IMA_POLICY_FILE='/etc/ima/ima-policy'

#=============================================================================
# Function Name: _fn_reset_ima_policy
# Description  : remove the generated area in the IMA policy file
#                If the generated area is incomplete, it will only remove the patterns.
#                If the generated area is complete, it will remove the lines between the patterns.
#=============================================================================
function _fn_reset_ima_policy()
{
    local start_pattern="$IMA_POLICY_GENERATED_PATTERN_START"
    local end_pattern="$IMA_POLICY_GENERATED_PATTERN_END"

    # find generated area
    # grep will output <line>:<content>, use cut to get the line number
    local start_line=$(grep -n "$start_pattern" "$IMA_POLICY_FILE" | cut -d: -f1)
    local end_line=$(grep -n "$end_pattern" "$IMA_POLICY_FILE" | cut -d: -f1)

    # if both patterns are not found, do nothing
    if [ -z "$start_line" ] && [ -z "$end_line" ]; then
        echo "No generated IMA policy found in $IMA_POLICY_FILE."
        return 0
    fi

    # incomplete generated area, reamove pattern only
    if [ -z "$start_line" ] || [ -z "$end_line" ]; then
        echo "Incomplete generated IMA policy found in $IMA_POLICY_FILE. Removing patterns only."
        sed -i "/$start_pattern/d" "$IMA_POLICY_FILE"
        sed -i "/$end_pattern/d" "$IMA_POLICY_FILE"
        return 0
    fi

    # remove lines between the patterns
    local lines_to_remove=$(sed -n "${start_line},${end_line}p" "$IMA_POLICY_FILE" | wc -l)
    if [ $lines_to_remove -gt 0 ]; then
        echo "Removing $lines_to_remove lines from IMA policy."
        sed -i "${start_line},${end_line}d" "$IMA_POLICY_FILE"
    fi
}

#=============================================================================
# Function Name: _fn_generate_from_config
# Description  : generate IMA policy from the configuration file
#                The configuration file should contain SELinux tags, one per line.
#                Each line will be checked if the tag is present in SELinux.
# Parameter    : config file path
#=============================================================================
function _fn_generate_from_config()
{
    local config_file="$1"
    # if last line is not empty, add a new line
    if [ -s $IMA_POLICY_FILE ] && [ "$(tail -c 1 $IMA_POLICY_FILE)" != "" ]; then
        echo "" >> $IMA_POLICY_FILE
    fi
    echo "$IMA_POLICY_GENERATED_PATTERN_START" >> $IMA_POLICY_FILE

    grep -v '^#' $config_file| grep -v '^$'| grep -Ev '^[[:space:]]+'| while read line
    do
        # skip empty lines and comments
        if [ -z "$line" ] || [[ "$line" =~ ^# ]]; then
            continue
        fi

        # check tag presence
        if [ -z "$(seinfo -t "$line" --flat)" ]; then
            fn_warn "SELinux tag $line is not present in selinux."
            continue
        fi

        # check if the tag is already present in the ima policy
        if grep -q "measure func=FILE_CHECK obj_type=$line" "$IMA_POLICY_FILE"; then
            echo "SELinux tag $line is already present in IMA policy."
            continue
        fi

        # append the line to the ima policy
        echo "measure func=FILE_CHECK obj_type=$line" >> $IMA_POLICY_FILE
        echo "Added SELinux tag $line to IMA policy."
    done
    unset line
    echo "$IMA_POLICY_GENERATED_PATTERN_END" >> $IMA_POLICY_FILE
}

#=============================================================================
# Function Name: _fn_activate_ima
# Description  : import ima config to ima
#=============================================================================
function _fn_activate_ima()
{
    cat $IMA_POLICY_FILE > /sys/kernel/security/ima/policy
    if [ $? -ne 0 ]; then
        fn_warn "Failed to activate IMA policy. The kernel may not enable CONFIG_IMA_WRITE_POLICY."
        fn_warn "You need reboot the system to apply the IMA policy."
    else
        echo "IMA policy activated successfully."
    fi
}

#=============================================================================
# Function Name: fn_ima_tool_main
# Description  : ima tool main function
# Parameter    : config file path
# Returns      : 0 on success, otherwise on fail
#=============================================================================
function fn_ima_tool_main()
{
    if [ ! $# -eq 1 ]; then
        echo "Please provide the configuration file path as an argument."
        exit 1
    fi
  
    local config_file="$1"
    # operator must be root
    if [ `id -u` -ne 0 ]; then
        echo "You must be logged in as root."
        exit 1
    fi

    if [ ! -f /sys/kernel/security/ima/policy ]; then
        echo "IMA is not enabled in the kernel. The IMA policy cannot be generated." >&2
        return 1
    fi

    if getenforce | grep -q "Disabled"; then
        fn_warn "SELinux is not enabled. IMA policy config will remain unchanged."
        return 1
    fi

    echo "Generating IMA policy..."
    if [ ! -f $IMA_POLICY_FILE ]; then
        echo "Creating IMA policy at $IMA_POLICY_FILE."
        touch $IMA_POLICY_FILE
        chmod 600 $IMA_POLICY_FILE
    else
        echo "Resetting existing IMA policy at $IMA_POLICY_FILE."
        _fn_reset_ima_policy
    fi

    # check if the config file exists
    if [ ! -f "$config_file" ]; then
        echo "ima configuration file $config_file does not exist." >&2
        return 1
    fi
    _fn_generate_from_config "$config_file"
    _fn_activate_ima
}

fn_ima_tool_main $@
